# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

cmake_minimum_required(VERSION 3.12)
project(msda LANGUAGES CXX CUDA)

find_package(CUDA REQUIRED)

set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/lib)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/bin)

set(arch ${CMAKE_HOST_SYSTEM_PROCESSOR})

if(DEFINED TRT_HOME)
  message("-- using TRT_HOME = ${TRT_HOME}")
  set(TENSORRT_INCLUDE ${TRT_HOME}/include)
  set(TENSORRT_LIB ${TRT_HOME}/lib)
else()
  if(CMAKE_SYSTEM_PROCESSOR MATCHES "^aarch64")
    set(TENSORRT_INCLUDE /usr/include/aarch64-linux-gnu/)
    set(TENSORRT_LIB /usr/lib/aarch64-linux-gnu/)
  elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^x86_64")
    set(TENSORRT_INCLUDE /usr/include/x86_64-linux-gnu/)
    set(TENSORRT_LIB /usr/lib/x86_64-linux-gnu/)
  else()
    message(FATAL_ERROR "Unsupported processor: ${CMAKE_SYSTEM_PROCESSOR}")
  endif()
endif()

set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} -lineinfo -gencode arch=compute_86,code=compute_86 -gencode arch=compute_87,code=compute_87)

include_directories(
  ${CUDA_INCLUDE_DIRS}
  ${TENSORRT_INCLUDE}
  common
)

link_directories(
  ${TENSORRT_LIB}
  build
)

file(GLOB_RECURSE PLUGIN_FILES csrc/*.cpp csrc/*.cu)
cuda_add_library(msda SHARED csrc/msda/msda.cu csrc/trt/msda_plugin.cpp csrc/trt/msda_plugin_v3.cpp)
target_include_directories(msda PRIVATE csrc/msda)
target_link_libraries(msda nvinfer nvinfer_plugin)

# followings are for torch extention. 
# you may remove it if you do not want to build torch extention
# for example on Orin

find_package(PythonInterp)
message("-- use python at " ${PYTHON_EXECUTABLE})
if(DEFINED PYTHON_EXECUTABLE)
  execute_process(COMMAND ${PYTHON_EXECUTABLE} -c "import pybind11; print(pybind11.get_cmake_dir())" OUTPUT_VARIABLE pybind11_DIR OUTPUT_STRIP_TRAILING_WHITESPACE)
  if(pybind11_DIR)
    message("-- use pybind11 at " ${pybind11_DIR})
    find_package(pybind11 REQUIRED)
  endif()
  execute_process(COMMAND ${PYTHON_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path + '/Torch')" OUTPUT_VARIABLE Torch_DIR OUTPUT_STRIP_TRAILING_WHITESPACE)
  if(Torch_DIR)
    set(Torch_LIB_DIRS ${Torch_DIR}/../../../lib)
    message("-- use torch at " ${Torch_DIR} " with additional libs locates at " ${Torch_LIB_DIRS})
    find_package(Torch REQUIRED)
    set(TORCH_CUDA_ARCH_LIST "7.5 8.6 8.7")
  endif()
endif()

if(Torch_FOUND)
  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
  cuda_add_library(msda_torch SHARED csrc/msda/msda.cu csrc/torch/msda_torch.cpp)
  target_include_directories(msda_torch PRIVATE csrc/msda)
  set_target_properties(msda_torch PROPERTIES CUDA_ARCHITECTURES "86,87")
  target_link_directories(msda_torch PRIVATE ${Torch_LIB_DIRS})
  target_link_libraries(msda_torch ${TORCH_LIBRARIES} torch_python)
  target_include_directories(msda_torch PRIVATE ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS})
endif()
